#!/usr/bin/env python

#it's better to call this script from cl_make_and_load_db.py because then a check is performed if all necessary files are present and some directories are created

#use $TMPDIR on cluster

import sys, os, os.path,time,subprocess, signal, mysql
from openeye.oechem import *

debug = False

if len(sys.argv) <9 and len(sys.argv) >10:
        print 'make_and_load_db.py <username> <password> <smiles file> <cluster: True|False> <update mysql db: True|False> <update mysql db with failed ligs: True|False> <project> <path to scripts>'
        sys.exit(1)

username = sys.argv[1]
password = sys.argv[2]
in_file = sys.argv[3]
cluster = sys.argv[4]
update_db = sys.argv[5]
update_failed = sys.argv[6]
project = sys.argv[7]
script_path = sys.argv[8]
if len(sys.argv)==10:
	smarts_file=sys.argv[9]
	print "smarts file given!"



#----------------------------
def reformat_mol2_file(format_file):
	#OElib now produces a format that is not compatible with mol2db -> change
	correct_file = open(format_file,'r')
        correct_file_text = correct_file.read()
        #print correct_file_text
        correct_file_text = correct_file_text.replace('NO_CHARGES', 'NO_CHARGES\n')
        correct_file.close()
        correct_file = open(format_file,'w')
        correct_file.write(correct_file_text)
        #print correct_file_text
        correct_file.close()
#----------------------------
def solv_remove(remove_id, solv_file):
	command = 'rm temp.solv'
        os.system(command)
        temp_solv = open('temp.solv', 'w')
	print solv_file
        org_solv_file = open(solv_file, 'r')
        for line in org_solv_file.readlines():
		#print line
                if line[0:3] == 'MFC':
                	title = line[0:12]
                	#print '----------- check ', title
                	if title in removed_id:
                		delete = True
                        else:
                        	delete = False
                if not delete:
                	temp_solv.write(line)
	temp_solv.close()
	org_solv_file.close()
#----------------------------




#check if input file exists -> is an issue with array jobs because they always expect consecetutive numbering which might not always be the case

if not os.path.exists(in_file):
	print in_file, 'does not exist'
	sys.exit()

if cluster == 'True':
	cluster = True
else:
	cluster = False


if update_db == 'True':
	update_db = True
else:
	update_db = False

if update_failed == 'True':
	update_failed = True
else:
	update_failed = False

if update_db or update_failed:

	blocked = True
	conn=mysql.connect2server(password, username,'purchasable')  
	cursor = conn.cursor ()
			


	#get table names

	command = "select stereoisomer_smiles from purchasable.projects where project = '" + project + "'"
	print command
	cursor.execute(command)
	results=cursor.fetchall()
	if len(results) == 0:
		print 'project does not exist'
		sys.exit(1)

	stereo_table = results[0][0]

	#get stereo_id,blob, error from table
	command = "show fields from " + stereo_table
	cursor.execute(command)
	results=cursor.fetchall()
	for result in results:
		print result[3], result[0]
		if result[3] == 'PRI':
			stereo_field = result[0]
		elif result[0].find('blob') <> -1 :
			blob_field = result[0]
		elif result[0].find('error') <> -1 :
			error_field = result[0]
	print stereo_field, blob_field, error_field
	error_field = blob_field


cur_dir = os.getcwd()


#make some needed directoryies
if not os.path.exists('omega_mult'):
	os.makedirs('omega_mult')
if not os.path.exists('solv'):
	os.makedirs('solv')
if not os.path.exists('errors'):
	os.makedirs('errors')	
if not os.path.exists('input_smiles'):
	os.makedirs('input_smiles')	


#make subdir
dir_name = in_file[:-4] + '_work_dir'
if cluster and not debug:
	jobid = os.getenv('LSB_JOBID')

	tmp_dir = '/jobdir/' + jobid + '/'
	dir_name = tmp_dir + '/' + dir_name
print dir_name

if os.path.exists(dir_name):
	print 'working dir already exists:', in_file
else:
	os.makedirs(dir_name)
	os.chdir(dir_name)
	os.makedirs('amsol2')
	os.makedirs('omega_mult')
	if cluster:
		command = 'cp ' + cur_dir + '/' + in_file + ' .'
	else:
		command = 'mv ../' + in_file + ' .'
	print command
	os.system(command)
	#smarts file
	if len(sys.argv)==10:
		command = 'cp ' + cur_dir + '/' + smarts_file + ' omega_mult'
		os.system(command)
	command = 'cp  ' + cur_dir + '/omega.* .'
	os.system(command)
	
	#omega
	if cluster:
		command = 'omega2 -in ' + in_file[:-4] + '.smi -out amsol2/' + in_file[:-4] + '_os.mol2 -param omega.single -log '+ in_file[:-4] + '_os.log'
	else:
		command = 'omega2 -in ' + in_file[:-4] + '.smi -out amsol2/' + in_file[:-4] + '_os.mol2 -param omega.single -log '+ in_file[:-4] + '_os.log'

	print command
	os.system(command)
	

	#omega mult
	if cluster:
		command = 'omega2 -in ' + 'amsol2/' + in_file[:-4] + '_os.mol2 -out omega_mult/' + in_file[:-4] + '_om.mol2 -param omega.mult -log '+ in_file[:-4] + '_om.log'
	else:
		command = 'omega2 -in ' + 'amsol2/'  + in_file[:-4] + '_os.mol2 -out omega_mult/' + in_file[:-4] + '_om.mol2 -param omega.mult -log '+ in_file[:-4] + '_om.log'
	print command
	os.system(command)
	

	#mark failed molecules, omega.single
	failed_file = open(in_file[:-4] + '_os.log', 'r')

	for line in failed_file.readlines():
		if line.find('Error') <> -1 or line.find('failed') <> -1:
			start = line.find('MFC')
			if start <> -1: #if "template failed" -> no MFC code in same line
				end = start + 12
				name = line[start:end]
				if update_db or update_failed:
					name = name.replace('MFC', '')
					command = 'update ' + stereo_table + ' set ' + error_field + " = 'omega.single' where " + stereo_field + ' =' + name	
					print command
					cursor.execute(command)
				else:
					print name, 'failed'
				
	failed_file.close()
	if update_db or update_failed:
		conn.commit() 
			
	#mark failed molecules, omega.mult
	#failed_file = open(in_file[:-4] + '_om.log', 'r')

	#for line in failed_file.readlines():
		#if line.find('Error') <> -1 or line.find('failed') <> -1:

			#start = line.find('MFC')
			#end = start + 12
			#name = line[start:end]
			#if update_db or update_failed:
				#name = name.replace('MFC', '')
				#command = 'update ' + stereo_table + ' set ' + error_field + " = 'omega.mult' where " + stereo_field + ' =' + name
				#print command
				#cursor.execute(command)
			#else:
				#print name, 'failed'
	#failed_file.close()  				
	#if update_db or update_failed:			
		#conn.commit()
		#cursor.close()
		#conn.close()

	#there is another file with failed molecules
	if os.path.exists('omega2.fail'):
		failed_file = open('omega2.fail', 'r')

		for line in failed_file.readlines():
			smiles,name = line.split(' ')
			name = name.strip()
			if update_db or update_failed:
				name = name.replace('MFC', '')
				command = 'update ' + stereo_table + ' set ' + error_field + " = 'omega' where " + stereo_field + ' =' + name
				print command
				cursor.execute(command)
			else:
				print name, 'failed'
			
		failed_file.close()  					
		if update_db or update_failed:
			conn.commit()
	
	#amsol
	command = script_path + 'amsol.py CYC cav'
	print command
	os.system(command)
	
	#mark failed molecules
	failed_amsol = os.listdir('./')
	#conn=mysql.connect2server(password, username)
	#cursor = conn.cursor ()
	for failed_file in failed_amsol:
		if failed_file[-3:] == 'CYC':
			failed_am = open(failed_file  + '/UpdateCharge.err', 'r')
			for code in failed_am.readlines():
				name = code.strip()
				if update_db or update_failed:
					name = name.replace('MFC', '')
					command = 'update ' + stereo_table + ' set ' + error_field + " = 'amsol' where " + stereo_field + ' =' + name
					print command
					cursor.execute(command)
				else:
					print name, 'failed'
			failed_am.close() 
	if update_db or update_failed:
		conn.commit()
		#cursor.close ()
		#conn.close
	
	#multio....
	os.chdir('omega_mult')
	command = script_path + 'superpos_conf_ensemble.py . .. False False'
	if len(sys.argv)==10:
		command=command+' '+smarts_file 
	print command
	os.system(command)
	
	reformat_mol2_file(in_file[:-4]+'_om_mult_rings.mol2')

	#mol2db
	os.system('cp ' + cur_dir + '/inhier .')
	omega_file = in_file[:-4] + '_om_mult_rings.mol2'
  	db_file = omega_file[:-5] + '.db'
  	command = 'ln -s ' + omega_file + ' temp.mol2'
  	print command
  	os.system(command)
	solv_file = '../' + in_file[:-4] + '_os_amsol_cav_CYC/' + in_file[:-4] + '_mult_rings_os_CYC.solv'
  	command = 'ln -s ' + solv_file +  ' temp.solv'
  	print command
  	os.system(command)

	moldb_command = ['mol2db', 'inhier']

	
	#sometimes mol2db hangs -> catch this
	mol2db_loop_finished = False #rerun mol2db until it finishes normally or no molecules are left
	removed_id = []
	attempts = 0
	while (not mol2db_loop_finished): 
		attempts = attempts + 1
		mol2db_log = in_file[:-4] + '_om_mult_rings.log'
		log_file = open(mol2db_log, 'w')
		print moldb_command
		proc = subprocess.Popen(moldb_command,stdout=log_file) #start process

		#let it try for ten minutes
		minutes = 1
		error = False
		mol2db_finished = False
		while (not mol2db_finished) and (minutes <= 6):
			time.sleep(10) #short sleep, enough for short jobs
			print proc.poll(), 'proc.poll()'
			if proc.poll() == 0: #job finished normally
				print 'finished normally'
				mol2db_finished = True
				error = False
			elif proc.poll() is not None: #The jobs has finished, but there was an error
				#do something
				#print proc.poll()
                         	print 'there was a problem with mol2db'
				mol2db_finished = True
				if proc.poll() == -11 and attempts > 1:
					error = False # there is a segmentation fault because no molecules are left => quit
				else:
					error = True

			elif proc.poll() is None: #job is still running 
				time.sleep(30) 
				minutes = minutes + 1
				print 'mol2db is still running'
		if not mol2db_finished:
		     #kill job	
		     print 'mol2db is hanging -> will now be killed'	
  		     print proc.pid
		     print signal.SIGUSR1
		     os.kill(proc.pid, signal.SIGUSR1)
		log_file.close()
		if error or not mol2db_finished: # mol2db did not run on all molecules => check log file
			log_file = open( mol2db_log , 'r')
			log_file_lines = log_file.readlines()
			log_file.close()
			#find last line which contains id number
			worked_id = ''
			for i in reversed(log_file_lines):
				#print i
				if len(i) == 74: #line with molecule id
					worked_id = 'MFC' + i.split()[0]
					#print worked_id
					break # found last working id

			if worked_id == '' and attempts > 0:
				#the first molecule caused the problem
				org_solv_file = open(solv_file, 'r')
				first_line = org_solv_file.readline()
				org_solv_file.close()
				failed = line[0:12]
				removed_id.append(failed)

			#clean mol2 file
			command = 'rm temp.mol2'
			os.system(command)
			#print omega_file
			omega_mols = oemolistream(omega_file)
			tmp_mols = oemolostream ('temp.mol2')
			mol = OEGraphMol()
			delete_following = False
			while OEReadMolecule(omega_mols, mol):
				title = mol.GetTitle()
				if title == worked_id:
					delete_following = True
				if title in removed_id:
					#skip
					continue #we do not want these molecules
				elif delete_following and title <> worked_id:
					delete_following = False #found the id number
					removed_id.append(title)
				else: 
					OEWriteMolecule(tmp_mols, mol)
					#print title


			print removed_id
			omega_mols.close()
			tmp_mols.close()
			reformat_mol2_file('temp.mol2')

			#adopt solv file
			solv_remove(removed_id, solv_file)
		if (not mol2db_finished) or error:
			mol2db_loop_finished = False # repeat mol2db after molis have been removed
		else:
			mol2db_loop_finished = True	
		
	print 'failed in mol2db', removed_id
	#update db
	if update_db or update_failed:
		for stereo_id in removed_id:
			stereo_id = stereo_id.replace('MFC','')
			stereo_id = stereo_id[0:-2] #last two digits are for ring systems
			command = 'update ' + stereo_table + ' set ' + error_field + " = 'mol2db' where " + stereo_field + ' =' + stereo_id 
			print command
			cursor.execute(command)	
			conn.commit()		



	command = 'mv temp.db ' + in_file[:-4] + '_om_mult_rings.db'
	os.system(command)

	#check that log file is okay
	log_file = open( in_file[:-4] + '_om_mult_rings.log' , 'r')
	log_file_text = log_file.readlines()
	log_file.close()
	b_fin = True
	b_err = False
	l_err = []
	for s_line in log_file_text:
		if 'Finished' in s_line:
			b_fin = False
		if 'error' in s_line and not 'Exiting' in s_line:
			b_err = True
			l_err.append(s_line.split()[-1])
	if b_err:
		b_fin = False		
		for s_name in l_err:
			if update_db or update_failed: #for sure only do something if update  is set
				command = 'update ' + stereo_table + ' set ' + error_field + " = 'failed' where " + stereo_field + ' =' + s_name[:-2]
				print command
				cursor.execute(command)
				conn.commit()		
			else:
				print 'MFC' + s_name,  'failed'
				
	if b_fin:		
		#did not work
		command = 'mv ../' + in_file  + ' ' + cur_dir + '/errors'
		print command
		os.system(command)
		
		
	else:
		#everything went fine
		if update_db:
			cursor.close ()  #close connection before opened new one in update_db section
			conn.close
		
		
		if debug:
			command = 'mv *mult*mol2* ' + cur_dir + '/omega_mult'
			print command
			os.system(command)
			command = 'mv ../*CYC/*mult*solv ' + cur_dir + '/solv'
			print command
			os.system(command)
		if update_db:
	
			#load file
			#header
			#command ='cat ../../header.txt *db > ../new.db'
			command = 'mv *db ../new.db'
			print command
			os.system(command)
			os.chdir('..')

			command = script_path + 'mysql_load_dockable.py ' + username + ' ' + password +  ' ' + stereo_table + ' ' + stereo_field + ' ' + blob_field +  ' ' + 'False'
			print command
			upload_success =  os.system(command)                   
			if upload_success == 0:
				print 'upload succesful'
				if not debug:
					if cluster:
						command = 'cp ' + cur_dir + '/' + in_file + ' ' + cur_dir + '/input_smiles'
					else:
						command = 'mv *smi ../input_smiles'
					print command
					os.system(command)
					os.chdir('..')
					if not cluster:					
						command = 'rm -rf ' + dir_name
						print command
						os.system(command)

			else:
				print 'could not upload molecules'
		elif not debug:
			command = 'mv *db ' + cur_dir + '/new_dbs'
			print command
			os.system(command)
			command = 'bzip2 ' + cur_dir + '/new_dbs/' +  in_file[:-4] + '_om_mult_rings.db'
			os.system(command)
			os.chdir('..')
			if not cluster:
				command = 'mv *smi ../input_smiles'
			else:
				command = 'mv ' + cur_dir + '/' + in_file + ' ' + cur_dir + '/input_smiles'
			print command
			os.system(command)
			os.chdir('..')
			if not cluster:
				print command
				os.system(command)

print 'make_and_load_dockabe.py done'

			

		

		
		
		
	
	
	

	
	
	

